from utils import *
import yaml
from modelFactory import *
from dataProcess import dataProcess
import numpy as np
import time
# import psutil
# from memory_profiler import profile

yaml_file = "/home/cyw/projects/malware_detected/configue.yaml"
with open(yaml_file, 'r') as f:
    config = yaml.safe_load(f)

class indicator:
    """
        指标相关的函数
    """
    def get_fpr(self, y_true, y_pred):
        # 这个指标需要确保没有出错，有的是二分类，有的是线性回归，可能会出现问题
        nbenign = (y_true == 0).sum()
        nfalse = (y_pred[y_true == 0] == 1).sum()
        if debugLog == "true":
            print("fpr：样本总数：{} 正常样本数：{} 错误分类的正常样本数：{}".format(len(y_true),nbenign,nfalse))
        return nfalse / float(nbenign)

    def get_fnr(self, y_true, y_pred):
        nMalware = (y_true == 1).sum()
        nfalse = (y_pred[y_true == 1] == 0).sum()
        if debugLog == "true":
            print("fnr: 样本总数：{} 恶意样本数：{} 错误分类的恶意样本数：{}".format(len(y_true),nMalware,nfalse))
        return nfalse / float(nMalware)
    
    def get_fn_fp_tn_tp(self,y_true,y_pred):
        nbenign = (y_true == 0).sum()
        nMalware = (y_true == 1).sum()
        fn = (y_pred[y_true == 1] == 0).sum()
        fp = (y_pred[y_true == 0] == 1).sum()
        tn = nbenign - fp
        tp = nMalware - fn
        print("fnr: 样本总数：{} 恶意样本数：{} 正常样本数{} fn：{} fp:{}".format(len(y_true),nMalware,nbenign,fn,fp))
        return [fn,fp,tn,tp]

class experimentRes:
    """
        模型的性能指标
    """
    def __init__(self):
        self.dataPro = dataProcess()
        self.inid = indicator()

    def find_threshold(self, y_true, y_pred, fpr_target):
        thresh = 0.0
        fpr = self.inid.get_fpr(y_true, y_pred > thresh)
        while fpr > fpr_target and thresh < 1.0:
            thresh += 0.0001
            fpr = self.inid.get_fpr(y_true, y_pred > thresh)
        return thresh, fpr
    
    def predict_PE_file(self, model, pePath):
        """
            输入pe样本路径，输出模型对其的预测值
        """
        putty_data = open(pePath, "rb").read()
        tar = self.dataPro.parse_pe2embedding(putty_data)
        return model.predict([tar])
    
    def differ_threshold(self,model,samplePath):
        """
            获得在不同pr下，模型的性能表现
        """
        X_test, y_test = read_vectorized_features(samplePath,"test")
        y_pred = model.predict(X_test)
        label = y_test
        if debugLog =="true":
            print("检测样本的数量为：{}".format(len(X_test)))
        tarFpr=[0,0.001,0.01,0.1,0.3,0.9,1]
        for tarValue in tarFpr:
            print("【indictor】fpr：{}%".format(tarValue*100))
            threshold, fpr = self.find_threshold(label, y_pred, tarValue)
            self.dataPro.gene_ectorized_form_by_path(samplePath,"test")
            XData, yData = read_vectorized_features(samplePath, "test")
            dataSavePath = os.path.join(samplePath,"res.jsonl")
            self.predict_PES_res(model,XData,yData,dataSavePath,threshold)

    # @profile
    def predict_PES_res(self,model,X_test,y_test,dataSavePath,threshold=0.5):
        """
            获得输入样本的检测率和误报率
        """
        startTime = time.time()
        y_pred = model.predict(X_test)
        label = y_test
        y=[]
        for i in range(len(y_pred)):
            if y_pred[i]>=threshold:
                y.append(1)
            else:
                y.append(0)
        y = np.array(y)
        if debugLog =="true":
            print("检测样本的数量为：{}".format(len(X_test)))
            print("特征向量长度：{}".format(len(X_test[0])))
            print("样本的特征向量：{}".format(X_test))
            print("样本的标签：{}".format(y_test))
            # print("样本的预测值：{}".format(y_pred))
            print("按照阈值划分后的值：{}".format(y))

        # print("ROC AUC:", roc_auc_score(y_test, y_pred))
        tempData = self.inid.get_fn_fp_tn_tp(label,y)
        fn = tempData[0]
        fp = tempData[1]
        tn = tempData[2]
        tp = tempData[3]
        print("Threshold: {:.3f}".format(threshold))
        print("误报率: {:.3f}%".format((fp/(fp+tn)) * 100))
        print("检测率: {:.3f}%".format((tp/(tp+fn))*100))
        print("1 - precision: {:.3f}%".format((1-tp/(tp+fp))*100))
        endTime = time.time()
        with open(dataSavePath,"a+") as file:
            for i in range(len(X_test)):
                resDict={}
                resDict["labelOrigin"]=str(int(label[i]))
                resDict["labelPred"]=str(y[i])
                resDict["preValue"]=str(round(y_pred[i],4))
                json_string = json.dumps(resDict)
                file.write(json_string+"\n")
        
    # @profile 
    def get_res_from_origin_res(self,dataPath, savePath, PEPaths,labels,model):
        """
            aim: 从指定的文件样本中获得模型的计算结果
            dataPath 数据集存放地址
            savaPath 中间结果存放地址
            PEPaths labels 数据集中文件夹名和文件夹对应的标签
            model 使用的机器学习模型
        """
        print("【1】正在提取样本特征")
        a = time.time()
        for i in range(len(PEPaths)):
            pePath = PEPaths[i]
            label = labels[i] 
            self.dataPro.parsePefromFile(dataPath, savePath, pePath, label)
        b = time.time()
        if debugLog == "true":
            print("预处理耗时：{}".format(b-a))
        print("【2】样本特征正在向量化")
        self.dataPro.gene_ectorized_form_by_path(savePath,"test")
        XData, yData = read_vectorized_features(savePath, "test")

        # process = psutil.Process()
        # info = process.memory_full_info()
        # print(f"模型计算前内存占用: {info.uss / 1024 ** 2:.2f} MB")
        print("【3】结果计算中")
        dataSavePath = os.path.join(savePath, "res.jsonl")
        a = time.time() 
        self.predict_PES_res(model,XData,yData,dataSavePath)
        b = time.time() 
        # process = psutil.Process()
        # info = process.memory_full_info()
        # print(f"模型计算后内存占用: {info.uss / 1024 ** 2:.2f} MB")
        if debugLog == "true":
            print("计算耗时：{}".format(b-a))
       
# @profile    
def main():
    print("---")
    model = modelFactory()
    eRes = experimentRes()
    a=time.time()
    dataPath = config["dataConfig"]["dataPath"]
    savePath = config["dataConfig"]["savePath"]
    eRes.get_res_from_origin_res(dataPath, savePath, ["benign","malware"], [0,1], model)
    # eRes.get_res_from_origin_res(dataPath, savePath, ["pe"], [1], model)
    b=time.time()
    print("样本的总的耗时{}".format(b-a))


    # ember数据集的测试样本
    # XData, yData = read_vectorized_features(r"/home/cyw/projects/malware_detected/emberData/ember2018","test")
    # a = time.time()
    # eRes.predict_PES_res(model,XData,yData,"/home/cyw/projects/malware_detected/emberData/ember2018/res.txt")
    # b = time.time()
    # print("ember数据集测试耗时：{}".format(b-a))


if __name__ == "__main__":
    main()
    # todo  优化目标,通过配置文件,切换模型和数据集,无需手动修改代码